#!/usr/bin/env python
# -*- encoding: utf-8 -*-
"""
    Sequence Labeling Model.
"""
import torch
import torch.nn as nn

from rnn import RNN
from crf import CRF

class LSTM_CRF(nn.Module):

    def __init__(self, **kwargs):
        """
        Args:
            rnn_input_dim:    int, input dims
            rnn_type:         str, options: ['rnn', 'lstm', 'gru']
            rnn_units:        int, rnn单元数, one direction
            num_rnn_layers:   int, 层数
            bi_flag:       bool, 是否双向, default is True
            use_crf:       bool, 是否使用crf层
            dropout_rate:  float, dropout rate
            average_batch: bool, 是否对batch的loss做平均
            use_cuda:      bool
        """
        super(LSTM_CRF, self).__init__()
        for k in kwargs:
            self.__setattr__(k, kwargs[k])
        self.forward_cat = None
        self.input_embed = None
        self.back_hidden = None
        self.win = 4

        assert self.win > 1 # suggest set 2~6
        # feature dropout
        self.dropout_feature = nn.Dropout(self.dropout_rate)

        # rnn layer
        self.rnn_layer = RNN(unit_type=self.rnn_type, input_dim=self.rnn_input_dim, hidden_units=self.rnn_units,
            num_layers=self.num_rnn_layers, bi_flag=self.bi_flag)

        # rnn dropout
        self.dropout_rnn = nn.Dropout(self.dropout_rate)

        # dense layer
        hidden_input_dim = self.rnn_input_dim + self.rnn_units
        self.hidden2tag = nn.Linear(hidden_input_dim, self.target_size+2)

        # crf layer loss
        if self.use_crf:
            self.loss_function = CRF(self.target_size+2, batch_first=True)
        else:
            self.loss_function = nn.CrossEntropyLoss(ignore_index=-100, size_average=False)

    def loss(self, feats, tags, mask):
        """
        Args:
            feats: size=(batch_size, seq_len, tag_size)
            tags:  size=(batch_size, seq_len)
            mask:  size=(batch_size, seq_len)
        """
        if not self.use_crf:
            batch_size, max_len = feats.size(0), feats.size(1)
            lstm_feats = feats.view(batch_size * max_len, -1)
            tags = tags.view(-1)
            loss_value = self.loss_function(lstm_feats, tags)
        else:
            # here we set tags pads
            batch_size, max_len = tags.shape
            xx = torch.zeros(batch_size, max_len, dtype=torch.long).cuda()
            pad_tags  = torch.where(tags >= 0, tags, xx)
            loss_value = -self.loss_function.forward(feats, pad_tags, mask)

        if self.average_batch:
            batch_size = feats.size(0)
            loss_value /= float(batch_size)

        return loss_value

    def forward(self, feat_input, states, sentlen=None):
        """
        Args:
             inputs: list
        """
        if states is None:
            # Set initial hidden and cell states
            states = (torch.zeros(self.num_rnn_layers, feat_input.size(0), self.rnn_units).cuda(),
                     torch.zeros(self.num_rnn_layers, feat_input.size(0), self.rnn_units).cuda())
            self.forward_cat = None
            self.input_embed = None
        # input shape
        batch_size = feat_input.size(0)
        feat_pass = self.dropout_feature(feat_input)
        rnn_outputs, states = self.rnn_layer(feat_pass, states) # [batch_size, seq_len, features]
        if self.forward_cat is None:
            self.forward_cat  = rnn_outputs
            self.input_embed  = feat_input
        else:
            self.forward_cat  = torch.cat((self.forward_cat, rnn_outputs), 1)
            self.input_embed  = torch.cat((self.input_embed, feat_input), 1)
                   
        self.back_hidden = torch.zeros(self.forward_cat.size(0), self.forward_cat.size(1), self.rnn_units).cuda()
        if sentlen is None:
            sentlen = [self.forward_cat.size(1)] * self.forward_cat.size(0)

        for ii, seq_len in enumerate(sentlen):
            if seq_len > self.win:
                self.back_hidden[ii, :(seq_len-self.win), :] = self.forward_cat[ii, self.win:seq_len, :] # 0:seq_len-self.win
                self.back_hidden[ii, (seq_len-self.win):(seq_len -1), :] = self.forward_cat[ii, seq_len-1, :] #.view(forward_cat.size(0), 1, forward_cat.size(2)).expand(-1, self.win-1, -1)
            else:
                if seq_len > 1:
                    self.back_hidden[ii, :(seq_len-1), :] = self.forward_cat[ii, seq_len-1, :] #.view(forward_cat.size(0), 1, forward_cat.size(2)).expand(-1, seq_len-1, -1)
        out_hidden = torch.cat((self.input_embed, self.back_hidden), -1)
        
        rnn_outputs = self.dropout_rnn(out_hidden.reshape(out_hidden.size(0)*out_hidden.size(1), out_hidden.size(2)))
        rnn_feats = self.hidden2tag(rnn_outputs)


        return rnn_feats.view(batch_size, self.forward_cat.size(1), -1), states

